#include "RtpHandler.h"

std::unordered_map<int, std::string> RtpHandler::NALU_TYPE = {{1, "P"}, {5, "IDR"}, {7, "SPS"}, {8, "PPS"}};

bool RtpHandler::RtpParser(char* data, int bytesReceived, int curBufferIndex, bool finishFlag, int& latestIdrIndex)
{
    start = 0;

    //当从起始位置解析时，判断上一帧是否有残留，nalu载荷额外处理
    if(needDataLength != 0)
    {
        if(!RtpParserRemain(data, bytesReceived, curBufferIndex, latestIdrIndex)) return true;
    }

    while (bytesReceived > start)
    {
        if(status == RTSP_INTERLEAVE)
        {
            if(CheckCut(data, RTSP_INTERLEAVE_LENGTH, bytesReceived)) return true;

            RTSPInterleave* interleave = (RTSPInterleave*)(data + start);
            curPayloadLength = ntohs(interleave->length);

            remainPayloadLength = curPayloadLength;
            start += RTSP_INTERLEAVE_LENGTH;
            if(interleave->channel == 0) status = RTP_HEADER;
            else
            {
                std::cout << "收到RTCP数据包"<< std::endl;
                status = RTCP_DATA;
            }
        }
        else if(status == RTP_HEADER)
        {
            if(CheckCut(data, RTP_HEADER_LENGTH, bytesReceived)) return true;

            RTPHeader* header = (RTPHeader*)(data + start);
            curTimestamp = ntohl(header->timestamp);
            curSeq = ntohs(header->sequenceNumber);

            remainPayloadLength -= RTP_HEADER_LENGTH;
            start += RTP_HEADER_LENGTH;
            status = NALU_HEADER;
        }
        else if(status == NALU_HEADER)
        {
            NaluHeader* header = (NaluHeader*)(data + start);
            int type = (int)header->type;

            if(type == 28)
            {
                status = NALU_FREG_HEADER;
            }
            else if(type >= 1 || type <= 9)
            {
                if(finishFlag)
                {
                    std::cout << "接收到足量数据，连接结束" << std::endl;
                    return false;
                }

                curType = type;
                if(RTP_LOG) std::cout << "获取NALU，index = " << curIndex++ << "，type = " << NALU_TYPE[curType]
                    << "，SEQ = " << curSeq << "，Timestamp = " << curTimestamp << std::endl;
                status = NALU_DATA;

                //单包
                endNode->packageType = 0;
                startNode = endNode;
            }
            else
            {
                std::cout << std::endl << "不支持的包类型：" << type << std::endl;
                return false;
            }

            remainPayloadLength -= NALU_HEADER_LENGTH;
            start += NALU_HEADER_LENGTH;
        }
        else if(status == NALU_FREG_HEADER)
        {
            NaluHeader* header = (NaluHeader*)(data + start);
            curType = (int)header->type;

            if(header->pre & 0x4)
            {
                if(finishFlag)
                {
                    std::cout << "接收到足量数据，连接结束" << std::endl;
                    return false;
                }
                startSeq = curSeq;

                //分包起始
                endNode->packageType = 1;
                startNode = endNode;
            }
            else if(header->pre & 0x2)
            {
                if(RTP_LOG) std::cout << "获取NALU，index = " << curIndex++ << "，type = " << NALU_TYPE[curType]
                  << "，SEQ = " << startSeq << '-'<< curSeq << "，Timestamp = " << curTimestamp << std::endl;

                //分包起始结束
                endNode->packageType = 2;
            }

            remainPayloadLength -= NALU_HEADER_LENGTH;
            start += NALU_HEADER_LENGTH;
            status = NALU_DATA;
        }
        else if(status == NALU_DATA)
        {
            endNode->SetStart(curBufferIndex + start, curType, curTimestamp);

            if(start + remainPayloadLength > bytesReceived)
            {
                remainDataLength = bytesReceived - start;
                needDataLength = remainPayloadLength - remainDataLength;
                return true;
            }

            start += remainPayloadLength;
            status = RTSP_INTERLEAVE;
            
            curNaluLength += endNode->SetEnd(curBufferIndex + start);
            if(endNode->packageType == 0 || endNode->packageType == 2)
            {
                startNode->naluLength = curNaluLength;
                curNaluLength = 0;

                if(startNode->naluType == 5) latestIdrIndex = startNode->nodeIndex;
            }
            endNode = endNode->next;
        }
        else if(status == RTCP_DATA)
        {
            if(start + remainPayloadLength > bytesReceived)
            {
                remainDataLength = bytesReceived - start;
                needDataLength = remainPayloadLength - remainDataLength;
                return true;
            }
            start += remainPayloadLength;
            status = RTSP_INTERLEAVE;
        }
    }

    return true;
}

bool RtpHandler::CheckCut(char* data, int pieceLength, int bytesReceived)
{
    if(start + pieceLength <= bytesReceived) return false;
    remainDataLength = bytesReceived - start;
    needDataLength = pieceLength - remainDataLength;
    memcpy(remainData, data + start, remainDataLength);
    return true;
}

bool RtpHandler::RtpParserRemain(char* data, int bytesReceived, int curBufferIndex, int& latestIdrIndex)
{
    //如果该数据包依旧无法使该阶段数据完整(不太可能出现)
    if(bytesReceived < needDataLength)
    {
        if(status != NALU_DATA && status != RTCP_DATA) memcpy(remainData + remainDataLength, data, bytesReceived);
        needDataLength = needDataLength - bytesReceived;
        remainDataLength += bytesReceived;
        return false;
    }
    
    if(status == RTSP_INTERLEAVE)
    {
        memcpy(remainData + remainDataLength, data, needDataLength);
        RTSPInterleave* interleave = (RTSPInterleave*)remainData;
        curPayloadLength = ntohs(interleave->length);

        remainPayloadLength = curPayloadLength;
        status = RTP_HEADER;
    }
    else if(status == RTP_HEADER)
    {
        memcpy(remainData + remainDataLength, data, needDataLength);
        RTPHeader* header = (RTPHeader*)remainData;
        curTimestamp = ntohl(header->timestamp);
        curSeq = ntohs(header->sequenceNumber);

        remainPayloadLength -= RTP_HEADER_LENGTH;
        status = NALU_HEADER;
    }
    else if(status == NALU_DATA)
    {
        status = RTSP_INTERLEAVE;
        curNaluLength += endNode->SetEnd(curBufferIndex + needDataLength);
        if(endNode->packageType == 0 || endNode->packageType == 2)
        {
            startNode->naluLength = curNaluLength;
            curNaluLength = 0;

            if(startNode->naluType == 5) latestIdrIndex = startNode->nodeIndex;
        }
        endNode = endNode->next;
    }
    else if(status == RTCP_DATA)
    {
        status = RTSP_INTERLEAVE;
    }

    start = needDataLength;
    needDataLength = 0;

    return true;
}

RtpHandler::RtpHandler(DataNode* _startNode)
{
    start = 0;
    status = RTSP_INTERLEAVE;
    curSeq = 0;
    startSeq = 0;
    curPayloadLength = 0;
    remainPayloadLength = 0;

    curIndex = 0;
    curTimestamp = 0;
    curType = -1;

    remainDataLength = 0;
    needDataLength = 0;

    startNode = _startNode;
    endNode = _startNode;
    curNaluLength = 0;
}

RtpHandler::~RtpHandler()
{
    std::cout << "RtpHandler已释放" << std::endl;
}